{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a40coVi7k4m-"
      },
      "outputs": [],
      "source": [
        "### This note book is designed to determine the number of parameters of STDC1\n",
        "### Source of code: The official GitHub repository of STDC (https://github.com/MichaelFan01/STDC-Seg)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install PrettyTable"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1g2mS3Nfk8BQ",
        "outputId": "db062cff-50e7-41d6-c668-6073be15527a"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: PrettyTable in /usr/local/lib/python3.10/dist-packages (0.7.2)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "### This code is taken from the STDC official GitHub repository. It determines the number of parameters of STDC1 backbone\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import init\n",
        "import math\n",
        "\n",
        "# for cunting number of parameters\n",
        "from prettytable import PrettyTable\n",
        "\n",
        "class ConvX(nn.Module):\n",
        "    def __init__(self, in_planes, out_planes, kernel=3, stride=1):\n",
        "        super(ConvX, self).__init__()\n",
        "        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel, stride=stride, padding=kernel//2, bias=False)\n",
        "        self.bn = nn.BatchNorm2d(out_planes)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.relu(self.bn(self.conv(x)))\n",
        "        return out\n",
        "\n",
        "\n",
        "class AddBottleneck(nn.Module):\n",
        "    def __init__(self, in_planes, out_planes, block_num=3, stride=1):\n",
        "        super(AddBottleneck, self).__init__()\n",
        "        assert block_num > 1, print(\"block number should be larger than 1.\")\n",
        "        self.conv_list = nn.ModuleList()\n",
        "        self.stride = stride\n",
        "        if stride == 2:\n",
        "            self.avd_layer = nn.Sequential(\n",
        "                nn.Conv2d(out_planes//2, out_planes//2, kernel_size=3, stride=2, padding=1, groups=out_planes//2, bias=False),\n",
        "                nn.BatchNorm2d(out_planes//2),\n",
        "            )\n",
        "            self.skip = nn.Sequential(\n",
        "                nn.Conv2d(in_planes, in_planes, kernel_size=3, stride=2, padding=1, groups=in_planes, bias=False),\n",
        "                nn.BatchNorm2d(in_planes),\n",
        "                nn.Conv2d(in_planes, out_planes, kernel_size=1, bias=False),\n",
        "                nn.BatchNorm2d(out_planes),\n",
        "            )\n",
        "            stride = 1\n",
        "\n",
        "        for idx in range(block_num):\n",
        "            if idx == 0:\n",
        "                self.conv_list.append(ConvX(in_planes, out_planes//2, kernel=1))\n",
        "            elif idx == 1 and block_num == 2:\n",
        "                self.conv_list.append(ConvX(out_planes//2, out_planes//2, stride=stride))\n",
        "            elif idx == 1 and block_num > 2:\n",
        "                self.conv_list.append(ConvX(out_planes//2, out_planes//4, stride=stride))\n",
        "            elif idx < block_num - 1:\n",
        "                self.conv_list.append(ConvX(out_planes//int(math.pow(2, idx)), out_planes//int(math.pow(2, idx+1))))\n",
        "            else:\n",
        "                self.conv_list.append(ConvX(out_planes//int(math.pow(2, idx)), out_planes//int(math.pow(2, idx))))\n",
        "\n",
        "    def forward(self, x):\n",
        "        out_list = []\n",
        "        out = x\n",
        "\n",
        "        for idx, conv in enumerate(self.conv_list):\n",
        "            if idx == 0 and self.stride == 2:\n",
        "                out = self.avd_layer(conv(out))\n",
        "            else:\n",
        "                out = conv(out)\n",
        "            out_list.append(out)\n",
        "\n",
        "        if self.stride == 2:\n",
        "            x = self.skip(x)\n",
        "\n",
        "        return torch.cat(out_list, dim=1) + x\n",
        "\n",
        "\n",
        "\n",
        "class CatBottleneck(nn.Module):\n",
        "    def __init__(self, in_planes, out_planes, block_num=3, stride=1):\n",
        "        super(CatBottleneck, self).__init__()\n",
        "        assert block_num > 1, print(\"block number should be larger than 1.\")\n",
        "        self.conv_list = nn.ModuleList()\n",
        "        self.stride = stride\n",
        "        if stride == 2:\n",
        "            self.avd_layer = nn.Sequential(\n",
        "                nn.Conv2d(out_planes//2, out_planes//2, kernel_size=3, stride=2, padding=1, groups=out_planes//2, bias=False),\n",
        "                nn.BatchNorm2d(out_planes//2),\n",
        "            )\n",
        "            self.skip = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)\n",
        "            stride = 1\n",
        "\n",
        "        for idx in range(block_num):\n",
        "            if idx == 0:\n",
        "                self.conv_list.append(ConvX(in_planes, out_planes//2, kernel=1))\n",
        "            elif idx == 1 and block_num == 2:\n",
        "                self.conv_list.append(ConvX(out_planes//2, out_planes//2, stride=stride))\n",
        "            elif idx == 1 and block_num > 2:\n",
        "                self.conv_list.append(ConvX(out_planes//2, out_planes//4, stride=stride))\n",
        "            elif idx < block_num - 1:\n",
        "                self.conv_list.append(ConvX(out_planes//int(math.pow(2, idx)), out_planes//int(math.pow(2, idx+1))))\n",
        "            else:\n",
        "                self.conv_list.append(ConvX(out_planes//int(math.pow(2, idx)), out_planes//int(math.pow(2, idx))))\n",
        "\n",
        "    def forward(self, x):\n",
        "        out_list = []\n",
        "        out1 = self.conv_list[0](x)\n",
        "\n",
        "        for idx, conv in enumerate(self.conv_list[1:]):\n",
        "            if idx == 0:\n",
        "                if self.stride == 2:\n",
        "                    out = conv(self.avd_layer(out1))\n",
        "                else:\n",
        "                    out = conv(out1)\n",
        "            else:\n",
        "                out = conv(out)\n",
        "            out_list.append(out)\n",
        "\n",
        "        if self.stride == 2:\n",
        "            out1 = self.skip(out1)\n",
        "        out_list.insert(0, out1)\n",
        "\n",
        "        out = torch.cat(out_list, dim=1)\n",
        "        return out\n",
        "\n",
        "#STDC2Net\n",
        "class STDCNet1446(nn.Module):\n",
        "    def __init__(self, base=64, layers=[4,5,3], block_num=4, type=\"cat\", num_classes=1000, dropout=0.20, pretrain_model='', use_conv_last=False):\n",
        "        super(STDCNet1446, self).__init__()\n",
        "        if type == \"cat\":\n",
        "            block = CatBottleneck\n",
        "        elif type == \"add\":\n",
        "            block = AddBottleneck\n",
        "        self.use_conv_last = use_conv_last\n",
        "        self.features = self._make_layers(base, layers, block_num, block)\n",
        "        self.conv_last = ConvX(base*16, max(1024, base*16), 1, 1)\n",
        "        self.gap = nn.AdaptiveAvgPool2d(1)\n",
        "        self.fc = nn.Linear(max(1024, base*16), max(1024, base*16), bias=False)\n",
        "        self.bn = nn.BatchNorm1d(max(1024, base*16))\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "        self.linear = nn.Linear(max(1024, base*16), num_classes, bias=False)\n",
        "\n",
        "        self.x2 = nn.Sequential(self.features[:1])\n",
        "        self.x4 = nn.Sequential(self.features[1:2])\n",
        "        self.x8 = nn.Sequential(self.features[2:6])\n",
        "        self.x16 = nn.Sequential(self.features[6:11])\n",
        "        self.x32 = nn.Sequential(self.features[11:])\n",
        "\n",
        "        if pretrain_model:\n",
        "            print('use pretrain model {}'.format(pretrain_model))\n",
        "            self.init_weight(pretrain_model)\n",
        "        else:\n",
        "            self.init_params()\n",
        "\n",
        "    def init_weight(self, pretrain_model):\n",
        "\n",
        "        state_dict = torch.load(pretrain_model)[\"state_dict\"]\n",
        "        self_state_dict = self.state_dict()\n",
        "        for k, v in state_dict.items():\n",
        "            self_state_dict.update({k: v})\n",
        "        self.load_state_dict(self_state_dict)\n",
        "\n",
        "    def init_params(self):\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.Conv2d):\n",
        "                init.kaiming_normal_(m.weight, mode='fan_out')\n",
        "                if m.bias is not None:\n",
        "                    init.constant_(m.bias, 0)\n",
        "            elif isinstance(m, nn.BatchNorm2d):\n",
        "                init.constant_(m.weight, 1)\n",
        "                init.constant_(m.bias, 0)\n",
        "            elif isinstance(m, nn.Linear):\n",
        "                init.normal_(m.weight, std=0.001)\n",
        "                if m.bias is not None:\n",
        "                    init.constant_(m.bias, 0)\n",
        "\n",
        "    def _make_layers(self, base, layers, block_num, block):\n",
        "        features = []\n",
        "        features += [ConvX(3, base//2, 3, 2)]\n",
        "        features += [ConvX(base//2, base, 3, 2)]\n",
        "\n",
        "        for i, layer in enumerate(layers):\n",
        "            for j in range(layer):\n",
        "                if i == 0 and j == 0:\n",
        "                    features.append(block(base, base*4, block_num, 2))\n",
        "                elif j == 0:\n",
        "                    features.append(block(base*int(math.pow(2,i+1)), base*int(math.pow(2,i+2)), block_num, 2))\n",
        "                else:\n",
        "                    features.append(block(base*int(math.pow(2,i+2)), base*int(math.pow(2,i+2)), block_num, 1))\n",
        "\n",
        "        return nn.Sequential(*features)\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat2 = self.x2(x)\n",
        "        feat4 = self.x4(feat2)\n",
        "        feat8 = self.x8(feat4)\n",
        "        feat16 = self.x16(feat8)\n",
        "        feat32 = self.x32(feat16)\n",
        "        if self.use_conv_last:\n",
        "           feat32 = self.conv_last(feat32)\n",
        "\n",
        "        return feat2, feat4, feat8, feat16, feat32\n",
        "\n",
        "    def forward_impl(self, x):\n",
        "        out = self.features(x)\n",
        "        out = self.conv_last(out).pow(2)\n",
        "        out = self.gap(out).flatten(1)\n",
        "        out = self.fc(out)\n",
        "        # out = self.bn(out)\n",
        "        out = self.relu(out)\n",
        "        # out = self.relu(self.bn(self.fc(out)))\n",
        "        out = self.dropout(out)\n",
        "        out = self.linear(out)\n",
        "        return out\n",
        "\n",
        "# STDC1Net\n",
        "class STDCNet813(nn.Module):\n",
        "    def __init__(self, base=64, layers=[2,2,2], block_num=4, type=\"cat\", num_classes=1000, dropout=0.20, pretrain_model='', use_conv_last=False):\n",
        "        super(STDCNet813, self).__init__()\n",
        "        if type == \"cat\":\n",
        "            block = CatBottleneck\n",
        "        elif type == \"add\":\n",
        "            block = AddBottleneck\n",
        "        self.use_conv_last = use_conv_last\n",
        "        self.features = self._make_layers(base, layers, block_num, block)\n",
        "        self.conv_last = ConvX(base*16, max(1024, base*16), 1, 1)\n",
        "        self.gap = nn.AdaptiveAvgPool2d(1)\n",
        "        self.fc = nn.Linear(max(1024, base*16), max(1024, base*16), bias=False)\n",
        "        self.bn = nn.BatchNorm1d(max(1024, base*16))\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "        self.linear = nn.Linear(max(1024, base*16), num_classes, bias=False)\n",
        "\n",
        "        self.x2 = nn.Sequential(self.features[:1])\n",
        "        self.x4 = nn.Sequential(self.features[1:2])\n",
        "        self.x8 = nn.Sequential(self.features[2:4])\n",
        "        self.x16 = nn.Sequential(self.features[4:6])\n",
        "        self.x32 = nn.Sequential(self.features[6:])\n",
        "\n",
        "        if pretrain_model:\n",
        "            print('use pretrain model {}'.format(pretrain_model))\n",
        "            self.init_weight(pretrain_model)\n",
        "        else:\n",
        "            self.init_params()\n",
        "\n",
        "    def init_weight(self, pretrain_model):\n",
        "\n",
        "        state_dict = torch.load(pretrain_model)[\"state_dict\"]\n",
        "        self_state_dict = self.state_dict()\n",
        "        for k, v in state_dict.items():\n",
        "            self_state_dict.update({k: v})\n",
        "        self.load_state_dict(self_state_dict)\n",
        "\n",
        "    def init_params(self):\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.Conv2d):\n",
        "                init.kaiming_normal_(m.weight, mode='fan_out')\n",
        "                if m.bias is not None:\n",
        "                    init.constant_(m.bias, 0)\n",
        "            elif isinstance(m, nn.BatchNorm2d):\n",
        "                init.constant_(m.weight, 1)\n",
        "                init.constant_(m.bias, 0)\n",
        "            elif isinstance(m, nn.Linear):\n",
        "                init.normal_(m.weight, std=0.001)\n",
        "                if m.bias is not None:\n",
        "                    init.constant_(m.bias, 0)\n",
        "\n",
        "    def _make_layers(self, base, layers, block_num, block):\n",
        "        features = []\n",
        "        features += [ConvX(3, base//2, 3, 2)]\n",
        "        features += [ConvX(base//2, base, 3, 2)]\n",
        "\n",
        "        for i, layer in enumerate(layers):\n",
        "            for j in range(layer):\n",
        "                if i == 0 and j == 0:\n",
        "                    features.append(block(base, base*4, block_num, 2))\n",
        "                elif j == 0:\n",
        "                    features.append(block(base*int(math.pow(2,i+1)), base*int(math.pow(2,i+2)), block_num, 2))\n",
        "                else:\n",
        "                    features.append(block(base*int(math.pow(2,i+2)), base*int(math.pow(2,i+2)), block_num, 1))\n",
        "\n",
        "        return nn.Sequential(*features)\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat2 = self.x2(x)\n",
        "        feat4 = self.x4(feat2)\n",
        "        feat8 = self.x8(feat4)\n",
        "        feat16 = self.x16(feat8)\n",
        "        feat32 = self.x32(feat16)\n",
        "        if self.use_conv_last:\n",
        "           feat32 = self.conv_last(feat32)\n",
        "        print(feat2.shape)\n",
        "        print(feat4.shape)\n",
        "        print(feat8.shape)\n",
        "        print(feat16.shape)\n",
        "        print(feat32.shape)\n",
        "\n",
        "        return feat2, feat4, feat8, feat16, feat32\n",
        "\n",
        "    def forward_impl(self, x):\n",
        "        out = self.features(x)\n",
        "        out = self.conv_last(out).pow(2)\n",
        "        out = self.gap(out).flatten(1)\n",
        "        out = self.fc(out)\n",
        "        # out = self.bn(out)\n",
        "        out = self.relu(out)\n",
        "        # out = self.relu(self.bn(self.fc(out)))\n",
        "        out = self.dropout(out)\n",
        "        out = self.linear(out)\n",
        "        return out\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    model = STDCNet813(num_classes=1000, dropout=0.00, block_num=4)\n",
        "    model.eval()\n",
        "    x = torch.randn(1,3,512,1024)\n",
        "    net = model(x)\n",
        "    #torch.save(STDC1.state_dict(), 'cat.pth')\n",
        "\n",
        "    #for printing the parameters\n",
        "    def count_parameters(model):\n",
        "        table = PrettyTable([\"Modules\", \"Parameters\"])\n",
        "        total_params = 0\n",
        "        for name, parameter in model.named_parameters():\n",
        "            if not parameter.requires_grad: continue\n",
        "            param = parameter.numel()\n",
        "            table.add_row([name, param])\n",
        "            total_params+=param\n",
        "        print(table)\n",
        "        print(f\"Total Trainable Params: {total_params}\")\n",
        "        return total_params\n",
        "\n",
        "    count_parameters(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AE-_zEubk8Ij",
        "outputId": "0a68d011-26db-4579-94e1-f0efb76fe1f8"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([1, 32, 256, 512])\n",
            "torch.Size([1, 64, 128, 256])\n",
            "torch.Size([1, 256, 64, 128])\n",
            "torch.Size([1, 512, 32, 64])\n",
            "torch.Size([1, 1024, 16, 32])\n",
            "+------------------------------------+------------+\n",
            "|              Modules               | Parameters |\n",
            "+------------------------------------+------------+\n",
            "|       features.0.conv.weight       |    864     |\n",
            "|        features.0.bn.weight        |     32     |\n",
            "|         features.0.bn.bias         |     32     |\n",
            "|       features.1.conv.weight       |   18432    |\n",
            "|        features.1.bn.weight        |     64     |\n",
            "|         features.1.bn.bias         |     64     |\n",
            "| features.2.conv_list.0.conv.weight |    8192    |\n",
            "|  features.2.conv_list.0.bn.weight  |    128     |\n",
            "|   features.2.conv_list.0.bn.bias   |    128     |\n",
            "| features.2.conv_list.1.conv.weight |   73728    |\n",
            "|  features.2.conv_list.1.bn.weight  |     64     |\n",
            "|   features.2.conv_list.1.bn.bias   |     64     |\n",
            "| features.2.conv_list.2.conv.weight |   18432    |\n",
            "|  features.2.conv_list.2.bn.weight  |     32     |\n",
            "|   features.2.conv_list.2.bn.bias   |     32     |\n",
            "| features.2.conv_list.3.conv.weight |    9216    |\n",
            "|  features.2.conv_list.3.bn.weight  |     32     |\n",
            "|   features.2.conv_list.3.bn.bias   |     32     |\n",
            "|   features.2.avd_layer.0.weight    |    1152    |\n",
            "|   features.2.avd_layer.1.weight    |    128     |\n",
            "|    features.2.avd_layer.1.bias     |    128     |\n",
            "| features.3.conv_list.0.conv.weight |   32768    |\n",
            "|  features.3.conv_list.0.bn.weight  |    128     |\n",
            "|   features.3.conv_list.0.bn.bias   |    128     |\n",
            "| features.3.conv_list.1.conv.weight |   73728    |\n",
            "|  features.3.conv_list.1.bn.weight  |     64     |\n",
            "|   features.3.conv_list.1.bn.bias   |     64     |\n",
            "| features.3.conv_list.2.conv.weight |   18432    |\n",
            "|  features.3.conv_list.2.bn.weight  |     32     |\n",
            "|   features.3.conv_list.2.bn.bias   |     32     |\n",
            "| features.3.conv_list.3.conv.weight |    9216    |\n",
            "|  features.3.conv_list.3.bn.weight  |     32     |\n",
            "|   features.3.conv_list.3.bn.bias   |     32     |\n",
            "| features.4.conv_list.0.conv.weight |   65536    |\n",
            "|  features.4.conv_list.0.bn.weight  |    256     |\n",
            "|   features.4.conv_list.0.bn.bias   |    256     |\n",
            "| features.4.conv_list.1.conv.weight |   294912   |\n",
            "|  features.4.conv_list.1.bn.weight  |    128     |\n",
            "|   features.4.conv_list.1.bn.bias   |    128     |\n",
            "| features.4.conv_list.2.conv.weight |   73728    |\n",
            "|  features.4.conv_list.2.bn.weight  |     64     |\n",
            "|   features.4.conv_list.2.bn.bias   |     64     |\n",
            "| features.4.conv_list.3.conv.weight |   36864    |\n",
            "|  features.4.conv_list.3.bn.weight  |     64     |\n",
            "|   features.4.conv_list.3.bn.bias   |     64     |\n",
            "|   features.4.avd_layer.0.weight    |    2304    |\n",
            "|   features.4.avd_layer.1.weight    |    256     |\n",
            "|    features.4.avd_layer.1.bias     |    256     |\n",
            "| features.5.conv_list.0.conv.weight |   131072   |\n",
            "|  features.5.conv_list.0.bn.weight  |    256     |\n",
            "|   features.5.conv_list.0.bn.bias   |    256     |\n",
            "| features.5.conv_list.1.conv.weight |   294912   |\n",
            "|  features.5.conv_list.1.bn.weight  |    128     |\n",
            "|   features.5.conv_list.1.bn.bias   |    128     |\n",
            "| features.5.conv_list.2.conv.weight |   73728    |\n",
            "|  features.5.conv_list.2.bn.weight  |     64     |\n",
            "|   features.5.conv_list.2.bn.bias   |     64     |\n",
            "| features.5.conv_list.3.conv.weight |   36864    |\n",
            "|  features.5.conv_list.3.bn.weight  |     64     |\n",
            "|   features.5.conv_list.3.bn.bias   |     64     |\n",
            "| features.6.conv_list.0.conv.weight |   262144   |\n",
            "|  features.6.conv_list.0.bn.weight  |    512     |\n",
            "|   features.6.conv_list.0.bn.bias   |    512     |\n",
            "| features.6.conv_list.1.conv.weight |  1179648   |\n",
            "|  features.6.conv_list.1.bn.weight  |    256     |\n",
            "|   features.6.conv_list.1.bn.bias   |    256     |\n",
            "| features.6.conv_list.2.conv.weight |   294912   |\n",
            "|  features.6.conv_list.2.bn.weight  |    128     |\n",
            "|   features.6.conv_list.2.bn.bias   |    128     |\n",
            "| features.6.conv_list.3.conv.weight |   147456   |\n",
            "|  features.6.conv_list.3.bn.weight  |    128     |\n",
            "|   features.6.conv_list.3.bn.bias   |    128     |\n",
            "|   features.6.avd_layer.0.weight    |    4608    |\n",
            "|   features.6.avd_layer.1.weight    |    512     |\n",
            "|    features.6.avd_layer.1.bias     |    512     |\n",
            "| features.7.conv_list.0.conv.weight |   524288   |\n",
            "|  features.7.conv_list.0.bn.weight  |    512     |\n",
            "|   features.7.conv_list.0.bn.bias   |    512     |\n",
            "| features.7.conv_list.1.conv.weight |  1179648   |\n",
            "|  features.7.conv_list.1.bn.weight  |    256     |\n",
            "|   features.7.conv_list.1.bn.bias   |    256     |\n",
            "| features.7.conv_list.2.conv.weight |   294912   |\n",
            "|  features.7.conv_list.2.bn.weight  |    128     |\n",
            "|   features.7.conv_list.2.bn.bias   |    128     |\n",
            "| features.7.conv_list.3.conv.weight |   147456   |\n",
            "|  features.7.conv_list.3.bn.weight  |    128     |\n",
            "|   features.7.conv_list.3.bn.bias   |    128     |\n",
            "|       conv_last.conv.weight        |  1048576   |\n",
            "|        conv_last.bn.weight         |    1024    |\n",
            "|         conv_last.bn.bias          |    1024    |\n",
            "|             fc.weight              |  1048576   |\n",
            "|             bn.weight              |    1024    |\n",
            "|              bn.bias               |    1024    |\n",
            "|           linear.weight            |  1024000   |\n",
            "+------------------------------------+------------+\n",
            "Total Trainable Params: 8443552\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "### This code is also taken from the STDC official GitHub repository. It determines the number of parameters of whole STDC1 end-to-end segmentation model.\n",
        "\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "\n",
        "\n",
        "#from nets.stdcnet import STDCNet1446, STDCNet813\n",
        "#from modules.bn import InPlaceABNSync as BatchNorm2d\n",
        "#BatchNorm2d = nn.BatchNorm2d\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as functional\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class ConvBNReLU(nn.Module):\n",
        "    def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):\n",
        "        super(ConvBNReLU, self).__init__()\n",
        "        self.conv = nn.Conv2d(in_chan,\n",
        "                out_chan,\n",
        "                kernel_size = ks,\n",
        "                stride = stride,\n",
        "                padding = padding,\n",
        "                bias = False)\n",
        "        # self.bn = BatchNorm2d(out_chan)\n",
        "        self.bn = nn.BatchNorm2d(out_chan)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.init_weight()\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv(x)\n",
        "        x = self.bn(x)\n",
        "        x = self.relu(x)\n",
        "        return x\n",
        "\n",
        "    def init_weight(self):\n",
        "        for ly in self.children():\n",
        "            if isinstance(ly, nn.Conv2d):\n",
        "                nn.init.kaiming_normal_(ly.weight, a=1)\n",
        "                if not ly.bias is None: nn.init.constant_(ly.bias, 0)\n",
        "\n",
        "\n",
        "class BiSeNetOutput(nn.Module):\n",
        "    def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):\n",
        "        super(BiSeNetOutput, self).__init__()\n",
        "        self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)\n",
        "        self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)\n",
        "        self.init_weight()\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv(x)\n",
        "        x = self.conv_out(x)\n",
        "        return x\n",
        "\n",
        "    def init_weight(self):\n",
        "        for ly in self.children():\n",
        "            if isinstance(ly, nn.Conv2d):\n",
        "                nn.init.kaiming_normal_(ly.weight, a=1)\n",
        "                if not ly.bias is None: nn.init.constant_(ly.bias, 0)\n",
        "\n",
        "    def get_params(self):\n",
        "        wd_params, nowd_params = [], []\n",
        "        for name, module in self.named_modules():\n",
        "            if isinstance(module, (nn.Linear, nn.Conv2d)):\n",
        "                wd_params.append(module.weight)\n",
        "                if not module.bias is None:\n",
        "                    nowd_params.append(module.bias)\n",
        "            elif isinstance(module, nn.BatchNorm2d):\n",
        "                nowd_params += list(module.parameters())\n",
        "        return wd_params, nowd_params\n",
        "\n",
        "\n",
        "class AttentionRefinementModule(nn.Module):\n",
        "    def __init__(self, in_chan, out_chan, *args, **kwargs):\n",
        "        super(AttentionRefinementModule, self).__init__()\n",
        "        self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)\n",
        "        self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)\n",
        "        # self.bn_atten = BatchNorm2d(out_chan)\n",
        "        self.bn_atten = nn.BatchNorm2d(out_chan)\n",
        "\n",
        "        self.sigmoid_atten = nn.Sigmoid()\n",
        "        self.init_weight()\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = self.conv(x)\n",
        "        atten = F.avg_pool2d(feat, feat.size()[2:])\n",
        "        atten = self.conv_atten(atten)\n",
        "        atten = self.bn_atten(atten)\n",
        "        atten = self.sigmoid_atten(atten)\n",
        "        out = torch.mul(feat, atten)\n",
        "        return out\n",
        "\n",
        "    def init_weight(self):\n",
        "        for ly in self.children():\n",
        "            if isinstance(ly, nn.Conv2d):\n",
        "                nn.init.kaiming_normal_(ly.weight, a=1)\n",
        "                if not ly.bias is None: nn.init.constant_(ly.bias, 0)\n",
        "\n",
        "\n",
        "class ContextPath(nn.Module):\n",
        "    def __init__(self, backbone='CatNetSmall', pretrain_model='', use_conv_last=False, *args, **kwargs):\n",
        "        super(ContextPath, self).__init__()\n",
        "\n",
        "        self.backbone_name = backbone\n",
        "        if backbone == 'STDCNet1446':\n",
        "            self.backbone = STDCNet1446(pretrain_model=pretrain_model, use_conv_last=use_conv_last)\n",
        "            self.arm16 = AttentionRefinementModule(512, 128)\n",
        "            inplanes = 1024\n",
        "            if use_conv_last:\n",
        "                inplanes = 1024\n",
        "            self.arm32 = AttentionRefinementModule(inplanes, 128)\n",
        "            self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)\n",
        "            self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)\n",
        "            self.conv_avg = ConvBNReLU(inplanes, 128, ks=1, stride=1, padding=0)\n",
        "\n",
        "        elif backbone == 'STDCNet813':\n",
        "            self.backbone = STDCNet813(pretrain_model=pretrain_model, use_conv_last=use_conv_last)\n",
        "            self.arm16 = AttentionRefinementModule(512, 128)\n",
        "            inplanes = 1024\n",
        "            if use_conv_last:\n",
        "                inplanes = 1024\n",
        "            self.arm32 = AttentionRefinementModule(inplanes, 128)\n",
        "            self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)\n",
        "            self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)\n",
        "            self.conv_avg = ConvBNReLU(inplanes, 128, ks=1, stride=1, padding=0)\n",
        "        else:\n",
        "            print(\"backbone is not in backbone lists\")\n",
        "            exit(0)\n",
        "\n",
        "        self.init_weight()\n",
        "\n",
        "    def forward(self, x):\n",
        "        H0, W0 = x.size()[2:]\n",
        "\n",
        "        feat2, feat4, feat8, feat16, feat32 = self.backbone(x)\n",
        "        H8, W8 = feat8.size()[2:]\n",
        "        H16, W16 = feat16.size()[2:]\n",
        "        H32, W32 = feat32.size()[2:]\n",
        "\n",
        "        avg = F.avg_pool2d(feat32, feat32.size()[2:])\n",
        "\n",
        "        avg = self.conv_avg(avg)\n",
        "        avg_up = F.interpolate(avg, (H32, W32), mode='nearest')\n",
        "\n",
        "        feat32_arm = self.arm32(feat32)\n",
        "        feat32_sum = feat32_arm + avg_up\n",
        "        feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')\n",
        "        feat32_up = self.conv_head32(feat32_up)\n",
        "\n",
        "        feat16_arm = self.arm16(feat16)\n",
        "        feat16_sum = feat16_arm + feat32_up\n",
        "        feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')\n",
        "        feat16_up = self.conv_head16(feat16_up)\n",
        "\n",
        "        return feat2, feat4, feat8, feat16, feat16_up, feat32_up # x8, x16\n",
        "\n",
        "    def init_weight(self):\n",
        "        for ly in self.children():\n",
        "            if isinstance(ly, nn.Conv2d):\n",
        "                nn.init.kaiming_normal_(ly.weight, a=1)\n",
        "                if not ly.bias is None: nn.init.constant_(ly.bias, 0)\n",
        "\n",
        "    def get_params(self):\n",
        "        wd_params, nowd_params = [], []\n",
        "        for name, module in self.named_modules():\n",
        "            if isinstance(module, (nn.Linear, nn.Conv2d)):\n",
        "                wd_params.append(module.weight)\n",
        "                if not module.bias is None:\n",
        "                    nowd_params.append(module.bias)\n",
        "            elif isinstance(module, nn.BatchNorm2d):\n",
        "                nowd_params += list(module.parameters())\n",
        "        return wd_params, nowd_params\n",
        "\n",
        "\n",
        "class FeatureFusionModule(nn.Module):\n",
        "    def __init__(self, in_chan, out_chan, *args, **kwargs):\n",
        "        super(FeatureFusionModule, self).__init__()\n",
        "        self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)\n",
        "        self.conv1 = nn.Conv2d(out_chan,\n",
        "                out_chan//4,\n",
        "                kernel_size = 1,\n",
        "                stride = 1,\n",
        "                padding = 0,\n",
        "                bias = False)\n",
        "        self.conv2 = nn.Conv2d(out_chan//4,\n",
        "                out_chan,\n",
        "                kernel_size = 1,\n",
        "                stride = 1,\n",
        "                padding = 0,\n",
        "                bias = False)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "        self.sigmoid = nn.Sigmoid()\n",
        "        self.init_weight()\n",
        "\n",
        "    def forward(self, fsp, fcp):\n",
        "        fcat = torch.cat([fsp, fcp], dim=1)\n",
        "        feat = self.convblk(fcat)\n",
        "        atten = F.avg_pool2d(feat, feat.size()[2:])\n",
        "        atten = self.conv1(atten)\n",
        "        atten = self.relu(atten)\n",
        "        atten = self.conv2(atten)\n",
        "        atten = self.sigmoid(atten)\n",
        "        feat_atten = torch.mul(feat, atten)\n",
        "        feat_out = feat_atten + feat\n",
        "        return feat_out\n",
        "\n",
        "    def init_weight(self):\n",
        "        for ly in self.children():\n",
        "            if isinstance(ly, nn.Conv2d):\n",
        "                nn.init.kaiming_normal_(ly.weight, a=1)\n",
        "                if not ly.bias is None: nn.init.constant_(ly.bias, 0)\n",
        "\n",
        "    def get_params(self):\n",
        "        wd_params, nowd_params = [], []\n",
        "        for name, module in self.named_modules():\n",
        "            if isinstance(module, (nn.Linear, nn.Conv2d)):\n",
        "                wd_params.append(module.weight)\n",
        "                if not module.bias is None:\n",
        "                    nowd_params.append(module.bias)\n",
        "            elif isinstance(module, nn.BatchNorm2d):\n",
        "                nowd_params += list(module.parameters())\n",
        "        return wd_params, nowd_params\n",
        "\n",
        "\n",
        "class BiSeNet(nn.Module):\n",
        "    def __init__(self, backbone, n_classes, pretrain_model='', use_boundary_2=False, use_boundary_4=False, use_boundary_8=False, use_boundary_16=False, use_conv_last=False, heat_map=False, *args, **kwargs):\n",
        "        super(BiSeNet, self).__init__()\n",
        "\n",
        "        self.use_boundary_2 = use_boundary_2\n",
        "        self.use_boundary_4 = use_boundary_4\n",
        "        self.use_boundary_8 = use_boundary_8\n",
        "        self.use_boundary_16 = use_boundary_16\n",
        "        # self.heat_map = heat_map\n",
        "        self.cp = ContextPath(backbone, pretrain_model, use_conv_last=use_conv_last)\n",
        "\n",
        "\n",
        "\n",
        "        if backbone == 'STDCNet1446':\n",
        "            conv_out_inplanes = 128\n",
        "            sp2_inplanes = 32\n",
        "            sp4_inplanes = 64\n",
        "            sp8_inplanes = 256\n",
        "            sp16_inplanes = 512\n",
        "            inplane = sp8_inplanes + conv_out_inplanes\n",
        "\n",
        "        elif backbone == 'STDCNet813':\n",
        "            conv_out_inplanes = 128\n",
        "            sp2_inplanes = 32\n",
        "            sp4_inplanes = 64\n",
        "            sp8_inplanes = 256\n",
        "            sp16_inplanes = 512\n",
        "            inplane = sp8_inplanes + conv_out_inplanes\n",
        "\n",
        "        else:\n",
        "            print(\"backbone is not in backbone lists\")\n",
        "            exit(0)\n",
        "\n",
        "        self.ffm = FeatureFusionModule(inplane, 256)\n",
        "        self.conv_out = BiSeNetOutput(256, 256, n_classes)\n",
        "        self.conv_out16 = BiSeNetOutput(conv_out_inplanes, 64, n_classes)\n",
        "        self.conv_out32 = BiSeNetOutput(conv_out_inplanes, 64, n_classes)\n",
        "\n",
        "        self.conv_out_sp16 = BiSeNetOutput(sp16_inplanes, 64, 1)\n",
        "\n",
        "        self.conv_out_sp8 = BiSeNetOutput(sp8_inplanes, 64, 1)\n",
        "        self.conv_out_sp4 = BiSeNetOutput(sp4_inplanes, 64, 1)\n",
        "        self.conv_out_sp2 = BiSeNetOutput(sp2_inplanes, 64, 1)\n",
        "        self.init_weight()\n",
        "\n",
        "    def forward(self, x):\n",
        "        H, W = x.size()[2:]\n",
        "\n",
        "        feat_res2, feat_res4, feat_res8, feat_res16, feat_cp8, feat_cp16 = self.cp(x)\n",
        "\n",
        "        feat_out_sp2 = self.conv_out_sp2(feat_res2)\n",
        "\n",
        "        feat_out_sp4 = self.conv_out_sp4(feat_res4)\n",
        "\n",
        "        feat_out_sp8 = self.conv_out_sp8(feat_res8)\n",
        "\n",
        "        feat_out_sp16 = self.conv_out_sp16(feat_res16)\n",
        "\n",
        "        feat_fuse = self.ffm(feat_res8, feat_cp8)\n",
        "\n",
        "        feat_out = self.conv_out(feat_fuse)\n",
        "        feat_out16 = self.conv_out16(feat_cp8)\n",
        "        feat_out32 = self.conv_out32(feat_cp16)\n",
        "\n",
        "        feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)\n",
        "        feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)\n",
        "        feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)\n",
        "\n",
        "\n",
        "        if self.use_boundary_2 and self.use_boundary_4 and self.use_boundary_8:\n",
        "            return feat_out, feat_out16, feat_out32, feat_out_sp2, feat_out_sp4, feat_out_sp8\n",
        "\n",
        "        if (not self.use_boundary_2) and self.use_boundary_4 and self.use_boundary_8:\n",
        "            return feat_out, feat_out16, feat_out32, feat_out_sp4, feat_out_sp8\n",
        "\n",
        "        if (not self.use_boundary_2) and (not self.use_boundary_4) and self.use_boundary_8:\n",
        "            return feat_out, feat_out16, feat_out32, feat_out_sp8\n",
        "\n",
        "        if (not self.use_boundary_2) and (not self.use_boundary_4) and (not self.use_boundary_8):\n",
        "            return feat_out, feat_out16, feat_out32\n",
        "\n",
        "    def init_weight(self):\n",
        "        for ly in self.children():\n",
        "            if isinstance(ly, nn.Conv2d):\n",
        "                nn.init.kaiming_normal_(ly.weight, a=1)\n",
        "                if not ly.bias is None: nn.init.constant_(ly.bias, 0)\n",
        "\n",
        "    def get_params(self):\n",
        "        wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []\n",
        "        for name, child in self.named_children():\n",
        "            child_wd_params, child_nowd_params = child.get_params()\n",
        "            if isinstance(child, (FeatureFusionModule, BiSeNetOutput)):\n",
        "                lr_mul_wd_params += child_wd_params\n",
        "                lr_mul_nowd_params += child_nowd_params\n",
        "            else:\n",
        "                wd_params += child_wd_params\n",
        "                nowd_params += child_nowd_params\n",
        "        return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "\n",
        "    net = BiSeNet('STDCNet813', 19)\n",
        "    #net.cuda()\n",
        "    net.eval()\n",
        "    #in_ten = torch.randn(1, 3, 512, 1024).cuda()\n",
        "    in_ten = torch.randn(1, 3, 512, 1024)\n",
        "    out, out16, out32 = net(in_ten)\n",
        "    print(out.shape)\n",
        "    torch.save(net.state_dict(), 'STDCNet813.pth')\n",
        "\n",
        "        #for printing the parameters\n",
        "    def count_parameters(model):\n",
        "        table = PrettyTable([\"Modules\", \"Parameters\"])\n",
        "        total_params = 0\n",
        "        for name, parameter in model.named_parameters():\n",
        "            if not parameter.requires_grad: continue\n",
        "            param = parameter.numel()\n",
        "            table.add_row([name, param])\n",
        "            total_params+=param\n",
        "        print(table)\n",
        "        print(f\"Total Trainable Params: {total_params}\")\n",
        "        return total_params\n",
        "\n",
        "    count_parameters(net)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QLnJzairk8MK",
        "outputId": "ba8347d6-ed68-46da-ffc5-2beb6669a625"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([1, 32, 256, 512])\n",
            "torch.Size([1, 64, 128, 256])\n",
            "torch.Size([1, 256, 64, 128])\n",
            "torch.Size([1, 512, 32, 64])\n",
            "torch.Size([1, 1024, 16, 32])\n",
            "torch.Size([1, 19, 512, 1024])\n",
            "+------------------------------------------------+------------+\n",
            "|                    Modules                     | Parameters |\n",
            "+------------------------------------------------+------------+\n",
            "|       cp.backbone.features.0.conv.weight       |    864     |\n",
            "|        cp.backbone.features.0.bn.weight        |     32     |\n",
            "|         cp.backbone.features.0.bn.bias         |     32     |\n",
            "|       cp.backbone.features.1.conv.weight       |   18432    |\n",
            "|        cp.backbone.features.1.bn.weight        |     64     |\n",
            "|         cp.backbone.features.1.bn.bias         |     64     |\n",
            "| cp.backbone.features.2.conv_list.0.conv.weight |    8192    |\n",
            "|  cp.backbone.features.2.conv_list.0.bn.weight  |    128     |\n",
            "|   cp.backbone.features.2.conv_list.0.bn.bias   |    128     |\n",
            "| cp.backbone.features.2.conv_list.1.conv.weight |   73728    |\n",
            "|  cp.backbone.features.2.conv_list.1.bn.weight  |     64     |\n",
            "|   cp.backbone.features.2.conv_list.1.bn.bias   |     64     |\n",
            "| cp.backbone.features.2.conv_list.2.conv.weight |   18432    |\n",
            "|  cp.backbone.features.2.conv_list.2.bn.weight  |     32     |\n",
            "|   cp.backbone.features.2.conv_list.2.bn.bias   |     32     |\n",
            "| cp.backbone.features.2.conv_list.3.conv.weight |    9216    |\n",
            "|  cp.backbone.features.2.conv_list.3.bn.weight  |     32     |\n",
            "|   cp.backbone.features.2.conv_list.3.bn.bias   |     32     |\n",
            "|   cp.backbone.features.2.avd_layer.0.weight    |    1152    |\n",
            "|   cp.backbone.features.2.avd_layer.1.weight    |    128     |\n",
            "|    cp.backbone.features.2.avd_layer.1.bias     |    128     |\n",
            "| cp.backbone.features.3.conv_list.0.conv.weight |   32768    |\n",
            "|  cp.backbone.features.3.conv_list.0.bn.weight  |    128     |\n",
            "|   cp.backbone.features.3.conv_list.0.bn.bias   |    128     |\n",
            "| cp.backbone.features.3.conv_list.1.conv.weight |   73728    |\n",
            "|  cp.backbone.features.3.conv_list.1.bn.weight  |     64     |\n",
            "|   cp.backbone.features.3.conv_list.1.bn.bias   |     64     |\n",
            "| cp.backbone.features.3.conv_list.2.conv.weight |   18432    |\n",
            "|  cp.backbone.features.3.conv_list.2.bn.weight  |     32     |\n",
            "|   cp.backbone.features.3.conv_list.2.bn.bias   |     32     |\n",
            "| cp.backbone.features.3.conv_list.3.conv.weight |    9216    |\n",
            "|  cp.backbone.features.3.conv_list.3.bn.weight  |     32     |\n",
            "|   cp.backbone.features.3.conv_list.3.bn.bias   |     32     |\n",
            "| cp.backbone.features.4.conv_list.0.conv.weight |   65536    |\n",
            "|  cp.backbone.features.4.conv_list.0.bn.weight  |    256     |\n",
            "|   cp.backbone.features.4.conv_list.0.bn.bias   |    256     |\n",
            "| cp.backbone.features.4.conv_list.1.conv.weight |   294912   |\n",
            "|  cp.backbone.features.4.conv_list.1.bn.weight  |    128     |\n",
            "|   cp.backbone.features.4.conv_list.1.bn.bias   |    128     |\n",
            "| cp.backbone.features.4.conv_list.2.conv.weight |   73728    |\n",
            "|  cp.backbone.features.4.conv_list.2.bn.weight  |     64     |\n",
            "|   cp.backbone.features.4.conv_list.2.bn.bias   |     64     |\n",
            "| cp.backbone.features.4.conv_list.3.conv.weight |   36864    |\n",
            "|  cp.backbone.features.4.conv_list.3.bn.weight  |     64     |\n",
            "|   cp.backbone.features.4.conv_list.3.bn.bias   |     64     |\n",
            "|   cp.backbone.features.4.avd_layer.0.weight    |    2304    |\n",
            "|   cp.backbone.features.4.avd_layer.1.weight    |    256     |\n",
            "|    cp.backbone.features.4.avd_layer.1.bias     |    256     |\n",
            "| cp.backbone.features.5.conv_list.0.conv.weight |   131072   |\n",
            "|  cp.backbone.features.5.conv_list.0.bn.weight  |    256     |\n",
            "|   cp.backbone.features.5.conv_list.0.bn.bias   |    256     |\n",
            "| cp.backbone.features.5.conv_list.1.conv.weight |   294912   |\n",
            "|  cp.backbone.features.5.conv_list.1.bn.weight  |    128     |\n",
            "|   cp.backbone.features.5.conv_list.1.bn.bias   |    128     |\n",
            "| cp.backbone.features.5.conv_list.2.conv.weight |   73728    |\n",
            "|  cp.backbone.features.5.conv_list.2.bn.weight  |     64     |\n",
            "|   cp.backbone.features.5.conv_list.2.bn.bias   |     64     |\n",
            "| cp.backbone.features.5.conv_list.3.conv.weight |   36864    |\n",
            "|  cp.backbone.features.5.conv_list.3.bn.weight  |     64     |\n",
            "|   cp.backbone.features.5.conv_list.3.bn.bias   |     64     |\n",
            "| cp.backbone.features.6.conv_list.0.conv.weight |   262144   |\n",
            "|  cp.backbone.features.6.conv_list.0.bn.weight  |    512     |\n",
            "|   cp.backbone.features.6.conv_list.0.bn.bias   |    512     |\n",
            "| cp.backbone.features.6.conv_list.1.conv.weight |  1179648   |\n",
            "|  cp.backbone.features.6.conv_list.1.bn.weight  |    256     |\n",
            "|   cp.backbone.features.6.conv_list.1.bn.bias   |    256     |\n",
            "| cp.backbone.features.6.conv_list.2.conv.weight |   294912   |\n",
            "|  cp.backbone.features.6.conv_list.2.bn.weight  |    128     |\n",
            "|   cp.backbone.features.6.conv_list.2.bn.bias   |    128     |\n",
            "| cp.backbone.features.6.conv_list.3.conv.weight |   147456   |\n",
            "|  cp.backbone.features.6.conv_list.3.bn.weight  |    128     |\n",
            "|   cp.backbone.features.6.conv_list.3.bn.bias   |    128     |\n",
            "|   cp.backbone.features.6.avd_layer.0.weight    |    4608    |\n",
            "|   cp.backbone.features.6.avd_layer.1.weight    |    512     |\n",
            "|    cp.backbone.features.6.avd_layer.1.bias     |    512     |\n",
            "| cp.backbone.features.7.conv_list.0.conv.weight |   524288   |\n",
            "|  cp.backbone.features.7.conv_list.0.bn.weight  |    512     |\n",
            "|   cp.backbone.features.7.conv_list.0.bn.bias   |    512     |\n",
            "| cp.backbone.features.7.conv_list.1.conv.weight |  1179648   |\n",
            "|  cp.backbone.features.7.conv_list.1.bn.weight  |    256     |\n",
            "|   cp.backbone.features.7.conv_list.1.bn.bias   |    256     |\n",
            "| cp.backbone.features.7.conv_list.2.conv.weight |   294912   |\n",
            "|  cp.backbone.features.7.conv_list.2.bn.weight  |    128     |\n",
            "|   cp.backbone.features.7.conv_list.2.bn.bias   |    128     |\n",
            "| cp.backbone.features.7.conv_list.3.conv.weight |   147456   |\n",
            "|  cp.backbone.features.7.conv_list.3.bn.weight  |    128     |\n",
            "|   cp.backbone.features.7.conv_list.3.bn.bias   |    128     |\n",
            "|       cp.backbone.conv_last.conv.weight        |  1048576   |\n",
            "|        cp.backbone.conv_last.bn.weight         |    1024    |\n",
            "|         cp.backbone.conv_last.bn.bias          |    1024    |\n",
            "|             cp.backbone.fc.weight              |  1048576   |\n",
            "|             cp.backbone.bn.weight              |    1024    |\n",
            "|              cp.backbone.bn.bias               |    1024    |\n",
            "|           cp.backbone.linear.weight            |  1024000   |\n",
            "|           cp.arm16.conv.conv.weight            |   589824   |\n",
            "|            cp.arm16.conv.bn.weight             |    128     |\n",
            "|             cp.arm16.conv.bn.bias              |    128     |\n",
            "|           cp.arm16.conv_atten.weight           |   16384    |\n",
            "|            cp.arm16.bn_atten.weight            |    128     |\n",
            "|             cp.arm16.bn_atten.bias             |    128     |\n",
            "|           cp.arm32.conv.conv.weight            |  1179648   |\n",
            "|            cp.arm32.conv.bn.weight             |    128     |\n",
            "|             cp.arm32.conv.bn.bias              |    128     |\n",
            "|           cp.arm32.conv_atten.weight           |   16384    |\n",
            "|            cp.arm32.bn_atten.weight            |    128     |\n",
            "|             cp.arm32.bn_atten.bias             |    128     |\n",
            "|           cp.conv_head32.conv.weight           |   147456   |\n",
            "|            cp.conv_head32.bn.weight            |    128     |\n",
            "|             cp.conv_head32.bn.bias             |    128     |\n",
            "|           cp.conv_head16.conv.weight           |   147456   |\n",
            "|            cp.conv_head16.bn.weight            |    128     |\n",
            "|             cp.conv_head16.bn.bias             |    128     |\n",
            "|            cp.conv_avg.conv.weight             |   131072   |\n",
            "|             cp.conv_avg.bn.weight              |    128     |\n",
            "|              cp.conv_avg.bn.bias               |    128     |\n",
            "|            ffm.convblk.conv.weight             |   98304    |\n",
            "|             ffm.convblk.bn.weight              |    256     |\n",
            "|              ffm.convblk.bn.bias               |    256     |\n",
            "|                ffm.conv1.weight                |   16384    |\n",
            "|                ffm.conv2.weight                |   16384    |\n",
            "|           conv_out.conv.conv.weight            |   589824   |\n",
            "|            conv_out.conv.bn.weight             |    256     |\n",
            "|             conv_out.conv.bn.bias              |    256     |\n",
            "|            conv_out.conv_out.weight            |    4864    |\n",
            "|          conv_out16.conv.conv.weight           |   73728    |\n",
            "|           conv_out16.conv.bn.weight            |     64     |\n",
            "|            conv_out16.conv.bn.bias             |     64     |\n",
            "|           conv_out16.conv_out.weight           |    1216    |\n",
            "|          conv_out32.conv.conv.weight           |   73728    |\n",
            "|           conv_out32.conv.bn.weight            |     64     |\n",
            "|            conv_out32.conv.bn.bias             |     64     |\n",
            "|           conv_out32.conv_out.weight           |    1216    |\n",
            "|         conv_out_sp16.conv.conv.weight         |   294912   |\n",
            "|          conv_out_sp16.conv.bn.weight          |     64     |\n",
            "|           conv_out_sp16.conv.bn.bias           |     64     |\n",
            "|         conv_out_sp16.conv_out.weight          |     64     |\n",
            "|         conv_out_sp8.conv.conv.weight          |   147456   |\n",
            "|          conv_out_sp8.conv.bn.weight           |     64     |\n",
            "|           conv_out_sp8.conv.bn.bias            |     64     |\n",
            "|          conv_out_sp8.conv_out.weight          |     64     |\n",
            "|         conv_out_sp4.conv.conv.weight          |   36864    |\n",
            "|          conv_out_sp4.conv.bn.weight           |     64     |\n",
            "|           conv_out_sp4.conv.bn.bias            |     64     |\n",
            "|          conv_out_sp4.conv_out.weight          |     64     |\n",
            "|         conv_out_sp2.conv.conv.weight          |   18432    |\n",
            "|          conv_out_sp2.conv.bn.weight           |     64     |\n",
            "|           conv_out_sp2.conv.bn.bias            |     64     |\n",
            "|          conv_out_sp2.conv_out.weight          |     64     |\n",
            "+------------------------------------------------+------------+\n",
            "Total Trainable Params: 12048928\n"
          ]
        }
      ]
    }
  ]
}